# In[ ]:
import os
import shutil
seed=6
import random
random.seed(seed)#主函数内设置随机数种子，对所有运行的子函数均起作用
import numpy as np
np.set_printoptions(suppress=True)
np.random.seed(seed)
os.environ['PYTHONHASHSEED']=str(seed)
# In[ ]:物理量常数
R0=6378135;
g0= 9.80665;
DTR=np.pi/180;
RTD=180/np.pi;

# In[ ]:造样本设置
timemaxmin=np.array([5000,200]);#最大/最小飞行时间
tranum=40000;#测试集样本量
trainnum=1000000;#训练集样本量
rpnum=2;#训练集中，对每一组随机的初始状态，生成剖面参数取不同值时的样本
sampsign=False;#是否重新造样本
modelload=False;
loadpoch=190000;
obnum=3;#初始状态量个数
actnum=9;#轨迹参数量个数
valuenum=10;#终端约束量个数,包括时间
# In[ ]:DLL初始化
if sampsign:
    alllist=['REENTRY.dll','REENTRY.exp','REENTRY.lib','REENTRY.pdb'];
    path=os.getcwd();
    parent = os.path.abspath(os.path.join(path, os.pardir));
    for filename in alllist:
        newpath=path+'\\'+filename;
        # if os.path.isfile(newpath)==False:
        oldpath=parent+'\\'+filename;
        shutil.copyfile(oldpath,newpath);

np.set_printoptions(suppress=True)

import ctypes
lib = ctypes.cdll.LoadLibrary(".\\REENTRY.dll")  
# 指定函数中的输入参数类型
#POINTER为返回的数组，转为numpy后一定要copy，防止随DLL改变
lib.Agent_stepForward.argtypes = (ctypes.POINTER(ctypes.c_double),)
lib.State_set.argtypes = (ctypes.POINTER(ctypes.c_double),)
lib.atreat.argtypes = (ctypes.c_bool,)
lib.amin_reset.argtypes = (ctypes.POINTER(ctypes.c_double),)
lib.amax_reset.argtypes = (ctypes.POINTER(ctypes.c_double),)
# 指定函数中的返回类型
#POINTER为返回的数组，转为numpy后一定要copy，防止随DLL改变
lib.Agent_stepForward.restype = ctypes.POINTER(ctypes.c_double)
lib.State_set.restype = ctypes.POINTER(ctypes.c_double)        
acttreatsign=True;#动作量超出范围的处理方式，False为硬限幅，True为用周期函数转化到区间的限幅
maxang=np.array([1,1],dtype=np.float32);#动作量最大值
minang=np.array([-1,-1],dtype=np.float32);#动作量最小值

lib.atreat(acttreatsign)
lib.amax_reset(np.ctypeslib.as_ctypes(maxang.astype(np.float64)))
lib.amin_reset(np.ctypeslib.as_ctypes(minang.astype(np.float64)))
lib.ainterval()
lib.tmax_reset(np.ctypeslib.as_ctypes(timemaxmin.astype(np.float64)));
# In[ ]:造样本
if sampsign:
    state0s=np.zeros((tranum,obnum));#初始状态
    params=np.zeros((tranum,actnum));#弹道参数
    Fstates=np.zeros((tranum,valuenum));#终端状态
    for trano in range(tranum):
        r = (50e3 + 10e3 * np.random.rand()) + R0;
        V = (5.0e3 + 0.5e3 * np.random.rand());
        gamma = -(2 * np.random.rand()) * DTR;
        state0=np.array([r,V,gamma]);
        state0s[trano,:]=state0;
        state=np.ctypeslib.as_array(lib.State_set(np.ctypeslib.as_ctypes(state0.astype(np.float64))), shape=[obnum]).astype(np.float32);
        
        dratio0 = np.random.rand();
        dratio_mid = np.random.rand();
        dratiof = np.random.rand();
        vratio = np.random.rand();
        ldratio0 = np.random.rand();
        ldratio_mid = np.random.rand();
        ldratiof = np.random.rand();
        Vf = 1500+1000*np.random.rand();
        reversetime=-500+1000*np.random.rand();
        param=np.array([dratio0,dratio_mid,dratiof,vratio,ldratio0,ldratio_mid,ldratiof,Vf,reversetime]);
        params[trano,:]=param;
        Fstate=np.ctypeslib.as_array(lib.Agent_stepForward(np.ctypeslib.as_ctypes(param.astype(np.float64))), shape=[valuenum]).astype(np.float32);
        Fstates[trano,:]=Fstate;    
        
    alldatax=(np.concatenate((state0s,params),axis=-1)).astype(np.float32);#样本集输入
    alldatay=(Fstates).astype(np.float32);#样本集输出

    state0s=np.zeros((trainnum,obnum));#初始状态
    params=np.zeros((trainnum,actnum));#弹道参数
    Fstates=np.zeros((trainnum,valuenum));#终端状态
    for trano in range(trainnum):
        if trano%rpnum==0:#重新生成初始状态
            r = (50e3 + 10e3 * np.random.rand()) + R0;
            V = (5.0e3 + 0.5e3 * np.random.rand());
            gamma = -(2 * np.random.rand()) * DTR;
        state0=np.array([r,V,gamma]);
        state0s[trano,:]=state0;
        state=np.ctypeslib.as_array(lib.State_set(np.ctypeslib.as_ctypes(state0.astype(np.float64))), shape=[obnum]).astype(np.float32);
        
        dratio0 = np.random.rand();
        dratio_mid = np.random.rand();
        dratiof = np.random.rand();
        vratio = np.random.rand();
        ldratio0 = np.random.rand();
        ldratio_mid = np.random.rand();
        ldratiof = np.random.rand();
        Vf = 1500+1000*np.random.rand();
        reversetime=-500+1000*np.random.rand();
        param=np.array([dratio0,dratio_mid,dratiof,vratio,ldratio0,ldratio_mid,ldratiof,Vf,reversetime]);
        params[trano,:]=param;
        Fstate=np.ctypeslib.as_array(lib.Agent_stepForward(np.ctypeslib.as_ctypes(param.astype(np.float64))), shape=[valuenum]).astype(np.float32);
        Fstates[trano,:]=Fstate;
        print(trano)
            
    datax=(np.concatenate((state0s,params),axis=-1)).astype(np.float32);#样本集输入
    datay=(Fstates).astype(np.float32);#样本集输出
    
    for s in ['datax','datay','alldatax','alldatay']:    
        np.savetxt(s+'.txt',np.array(eval(s)),fmt='%.5g')
else:
    for s in ['datax','datay','alldatax','alldatay']:    
        exec(s+'=(np.loadtxt(\''+s+'.txt\''+',ndmin=2'+')).astype(np.float32)')
       
def transtate(oristate):#把动力学输出的终端状态取所需的，单位一致
    #输出为角度制    
    if oristate.ndim==2:   
        fstate=np.concatenate((oristate[:,[1,2]]*RTD,oristate[:,[3]]),axis=-1);   
    else:
        fstate=np.concatenate((oristate[[1,2]]*RTD,oristate[[3]]),axis=-1);
        fstate=np.squeeze(fstate)
    return fstate

datay=transtate(datay);
alldatay=transtate(alldatay);
# datay=np.concatenate((alldatay,datay),axis=0);#训练时用所有样本
# datax=np.concatenate((alldatax,datax),axis=0);#训练时用所有样本
# In[ ]:数据处理dataset
import torch
torch.set_default_dtype(torch.float64)
from sklearn.preprocessing import MinMaxScaler;
yscaler=MinMaxScaler();

datax=datax[0:40000,:];
datay=datay[0:40000,:];
yscaler.fit(datay);

alldatax=datax[0:10000,:];
alldatay=datay[0:10000,:];
MMdatay=yscaler.fit_transform(datay);
MMalldatay=yscaler.fit_transform(alldatay);
dataset={'train_input':torch.tensor(datax,dtype=torch.float64),\
         'test_input':torch.tensor(alldatax,dtype=torch.float64),\
         'train_label':torch.tensor(MMdatay,dtype=torch.float64),\
         'test_label':torch.tensor(MMalldatay,dtype=torch.float64)}
inpnum=datax.shape[-1]
outnum=datay.shape[-1]
# In[ ]:KAN网络训练
miter=100;
steps=600;
from kan import *
from kan.utils import ex_round

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(device)

# create a KAN: k-order of spline, grid-grid intervals,grid_eps =0 adaptive grid
model = KAN(width=[inpnum,64,16,outnum], grid=10, k=3, seed=42, grid_eps =0.0,device=device)

ndatay0=model(dataset['train_input']);

model.update_grid_from_samples(dataset['train_input'])

with open('netloss.txt','w')as f:
    pass
for ii in range(1,miter+1):#自适应更新样条grid
    model.fit(dataset, opt="LBFGS", steps=steps,stop_grid_update_step=steps+1,batch=512,lamb=0.0001,lamb_entropy=0.001);
    savename='model'+str(ii)
    model.saveckpt('./'+savename)    
    ndatay=yscaler.inverse_transform(model(dataset['test_input']).detach().numpy());#输出反归一化
    wholeloss=np.mean(np.square(ndatay-alldatay),axis=(0));
    maxloss=np.max(np.abs(ndatay-alldatay),axis=(0))
    wholepoch=ii*steps
    with open('netloss.txt','a')as f:
        np.savetxt(f,[np.concatenate((np.array([wholepoch]),wholeloss,maxloss),axis=-1)],fmt='%.5g');#策略训练后记录critic网络估计actor网络动作值函数的过程量   
    print('iter'+str(wholepoch)+'test'+str(wholeloss)+';MAX'+str(maxloss));
# model2=KAN.loadckpt('./'+savename)#读取网络
# model.plot()
# In[ ]:
# lib = ['x','x^2','x^3','x^4','exp','log','sqrt','tanh','sin','abs']
# model.auto_symbolic(lib=lib)
# ex_round(model.symbolic_formula(),4)
